import torch
import torchvision
import torchvision.transforms as transforms
import torch.optim as optim
import torch.nn.functional as F
import random
import numpy as np

def determine_result(seed=114514):
    random.seed(seed) # 固定 Python 的随机种子
    np.random.seed(seed) # 固定 NumPy 的随机种子
    torch.manual_seed(seed) # 固定 CPU 上的 PyTorch 随机种子
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed) # 固定 GPU 上的 PyTorch 随机种子
        torch.cuda.manual_seed_all(seed) # 如果使用多个 GPU，则固定所有 GPU 的随机种子
        torch.backends.cudnn.deterministic = True # 确保 cuDNN 使用确定性的算法
        torch.backends.cudnn.benchmark = False # 关闭 cuDNN 的 benchmark 模式

def save_checkpoint(checkpoint, filename):
    print("=> Saving Checkpoint...")
    torch.save(checkpoint, filename)

def load_checkpoint(filepath):
    print("=> Loading Checkpoint...")
    ckpt = torch.load(filepath, weights_only=True)
    return ckpt['state_dict'], ckpt['optimizer']

def accuracy(output, target):
    return output.argmax(dim=1).eq(target).float().mean()

def evaluate(model, test_loader, device):
    model.eval()
    with torch.no_grad():
        acc_list = []
        loss_list = []
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            loss = F.cross_entropy(output, target)
            acc = accuracy(output, target)
            acc_list.append(acc.item())
            loss_list.append(loss.item())
        print(f"accuracy: {sum(acc_list) / len(acc_list)}, loss: {sum(loss_list) / len(loss_list)}")


def train(model, optimizer, train_loader, device):
    model.train()
    acc_list = []
    loss_list = []
    for data, target in train_loader:
        data, target = data.to(device), target.to(device)
        # forward
        scores = model(data)
        loss = F.cross_entropy(scores, target)
        acc = accuracy(scores, target)

        # backward
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        acc_list.append(acc.item())
        loss_list.append(loss.item())

    print(f"accuracy: {sum(acc_list) / len(acc_list)}, loss: {sum(loss_list) / len(loss_list)}")


def main():
    datasets_path = "/mnt/DISK/xjk/pytorch_learning/datasets/"
    checkpoints_path = "/mnt/DISK/xjk/pytorch_learning/checkpoints/"
    checkpoint_name = "checkpoint.pth.tar.pytorch_learning.loadsave"

    # load test data
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    test_data = torchvision.datasets.CIFAR10(root=datasets_path, train=False, transform=transforms.ToTensor(), download=True)
    test_loader = torch.utils.data.DataLoader(test_data, batch_size=64, shuffle=False)

    train_data = torchvision.datasets.CIFAR10(root=datasets_path, train=True, transform=transforms.ToTensor(), download=True)
    train_loader = torch.utils.data.DataLoader(train_data, batch_size=64, shuffle=True)

    model = torchvision.models.resnet18(weights=torchvision.models.ResNet18_Weights.DEFAULT)
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
    model.fc = torch.nn.Linear(512, 10)
    # summary(model, (1, 3, 32, 32))
    model.to(device)

    # train
    train(model, optimizer, train_loader, device)
    # evaluate
    evaluate(model, test_loader, device)

    # 一般情况下也会保存 optimizer 用于继续训练来进行微调
    checkpoint = {'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict()}
    save_checkpoint(checkpoint, checkpoints_path + checkpoint_name)

    model = torchvision.models.resnet18()
    model.fc = torch.nn.Linear(512, 10)
    model.to(device)
    # evaluate
    evaluate(model, test_loader, device)

    checkpoint, _ = load_checkpoint(checkpoints_path + checkpoint_name)
    model.load_state_dict(checkpoint)

    # evaluate again
    evaluate(model, test_loader, device)

if __name__ == "__main__":
    main()
